[None][perf] Reordering torch.compile and cache-dit call by BrianLi23 · Pull Request #16508 · NVIDIA/TensorRT-LLM

BrianLi23 · 2026-07-16T16:48:38Z

Summary

Move cache-acceleration enablement (TeaCache / Cache-DiT) out of each pipeline's post_load_weights into PipelineLoader.load, so it runs after torch.compile; Cache-DiT captures block-module references at enable time, so enabling it first left it driving stale eager blocks and torch.compile contributed nothing.
Centralize the call in the loader for all pipelines (Flux, Flux2, Wan T2V/I2V, LTX-2), leaving TeaCache coefficient resolution/validation in post_load_weights.
Skip cache_dit's inspect-based forward-pattern and output-count checks when transformer blocks are torch._dynamo.OptimizedModule wrappers, since compiled (*args, **kwargs) forwards can't match the declared patterns.
Update TeaCache unit tests to invoke _setup_cache_acceleration explicitly, mirroring the new loader-driven ordering.

Test Coverage

Model	Mean Before (s)	Mean After (s)	Std Before (s)	Std After (s)
Flux2	3.4827	3.2402	0.0183	0.0178
Qwen-Image	1.7481	1.5432	0.0093	0.0048
WAN2.2	15.8647	11.8930	0.0278	0.0112
LTX2	2.7206	2.7227	0.0148	0.0027

Note: the Qwen-Image results were measured with the Cache-DiT support from #16339 applied on top of this change, since Qwen-Image caching is not yet in main.

Flux2 Before

Flux2 After

Qwen-Image Before

Qwen-Image After

WAN2.2 Before

wan22-unpatched.mp4

WAN2.2 After
https://github.com/user-attachments/assets/56498c2b-f5e8-4eed-883e-b563257d6d05

Follow up

LTX2 shows ~2.03s denoise for pre/post changes. Cache-DiT works fine in both runs (6/20 steps cached), but torch.compile seems to be broken on LTX2 as transformer_ltx2.py:1033 (v_kv_list[idx]) guards on the static block index which causes dynamo to hit the 128-recompile limit which falls back to eager. Follow-up: fix the LTX2 index guard so compile survives; until then LTX2 gets zero compile benefit even without caching.
Reapply these changes on top of [None][feat] Qwen-Image TeaCache/Cache-DiT Support #16339 to reflect changes on Qwen Image.

Summary by CodeRabbit

Bug Fixes
- Improved TeaCache and Cache-DiT compatibility with compiled transformer blocks.
- Ensured cache acceleration is initialized in the correct order after model compilation.
- Improved validation handling for compiled blocks and multi-transformer pipelines.
Tests
- Updated TeaCache coverage for LTX-2 and Wan 2.2 video pipelines.

Signed-off-by: BrianLi23 <briannli233@gmail.com>

BrianLi23 · 2026-07-16T17:26:00Z

/bot help

github-actions · 2026-07-16T17:26:11Z

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental) --high-priority]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Supports wildcard * for pattern matching (e.g., "*PerfSanity*" matches all stages containing PerfSanity). Examples: "A10-PyTorch-1, xxx", "PerfSanity". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Supports wildcard * for pattern matching. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx", --extra-stage "Post-Merge".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

Signed-off-by: BrianLi23 <briannli233@gmail.com>

BrianLi23 · 2026-07-16T20:12:33Z

cc: @zhenhuaw-me

coderabbitai · 2026-07-16T21:35:01Z

📝 Walkthrough

Walkthrough

Changes

Cache acceleration flow

Layer / File(s)	Summary
Defer cache setup from pipeline post-load hooks `tensorrt_llm/_torch/visual_gen/models/{flux,ltx2,wan}/*`	Pipeline post-load methods retain TeaCache coefficient registration or validation while removing direct cache acceleration setup.
Enable acceleration after compilation `tensorrt_llm/_torch/visual_gen/pipeline_loader.py`	`PipelineLoader.load()` invokes `_setup_cache_acceleration()` after `torch.compile()` when a transformer is present.
Adapt compiled-block validation `tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py`, `tests/unittest/_torch/visual_gen/test_teacache.py`	FLUX Cache-DiT detects compiled blocks and disables incompatible forward-pattern checks; TeaCache tests explicitly invoke the relocated setup step.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PipelineLoader
  participant TorchCompile
  participant Pipeline
  participant CacheDiTAdapter
  PipelineLoader->>TorchCompile: compile pipeline when enabled
  TorchCompile-->>PipelineLoader: return compiled transformer blocks
  PipelineLoader->>Pipeline: _setup_cache_acceleration()
  Pipeline->>CacheDiTAdapter: configure compiled blocks
  CacheDiTAdapter-->>Pipeline: disable forward-pattern checks

Suggested reviewers: o-stoner

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 12.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately reflects the main change: reordering torch.compile with cache-DiT/TeaCache setup.
Description check	✅ Passed	The description includes a clear summary and test coverage table; the missing PR checklist section is non-critical.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/visual_gen/test_teacache.py`:
- Line 365: Add loader-level test coverage for the compiled Cache-DiT path by
exercising PipelineLoader.load() rather than only
BasePipeline._setup_cache_acceleration(). Assert that torch compilation occurs
before cache acceleration setup, or add an equivalent compiled Cache-DiT case
while preserving the existing direct setup tests.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: cbb8b9cb-1439-40e5-9bdd-801e5bfbeb41

📥 Commits

Reviewing files that changed from the base of the PR and between a2595c0 and d562d09.

📒 Files selected for processing (8)

tensorrt_llm/_torch/visual_gen/cache/cache_dit_enablers.py
tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py
tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py
tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py
tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py
tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py
tensorrt_llm/_torch/visual_gen/pipeline_loader.py
tests/unittest/_torch/visual_gen/test_teacache.py

Signed-off-by: BrianLi23 <briannli233@gmail.com>

zhenhuaw-me · 2026-07-17T00:35:46Z

@luyiyun1021 could help to take a look?

zhenhuaw-me

LGTM. @luyiyun1021 @o-stoner could you also take a look? Thanks!

zhenhuaw-me · 2026-07-17T01:36:51Z

/bot run

tensorrt-cicd · 2026-07-17T01:42:49Z

PR_Github #59821 [ run ] triggered by Bot. Commit: d562d09 Link to invocation

tensorrt-cicd · 2026-07-17T02:04:03Z

PR_Github #59821 [ run ] completed with state FAILURE. Commit: d562d09
/LLM/main/L0_MergeRequest_PR pipeline #48230 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

luyiyun1021 · 2026-07-21T03:12:15Z

The new SUPPORTED_CACHE_BACKENDS gate logs a warning and returns when the configured backend isn't in the pipeline's tuple. Before this PR, LTX2TwoStagesPipeline + cache_backend=cache_dit inherited LTX2Pipeline.post_load_weights's cache_dit branch and raised ValueError("Cache-DiT: no enabler registered for pipeline class 'LTX2TwoStagesPipeline'..."). Now the run proceeds at full denoise cost with one log warning. Do you think this works fine? @zhenhuaw-me

luyiyun1021 · 2026-07-21T03:13:49Z

Leave LTX enable cache dit with torch compile as a followup

zhenhuaw-me · 2026-07-21T06:32:19Z

The new SUPPORTED_CACHE_BACKENDS gate logs a warning and returns when the configured backend isn't in the pipeline's tuple. Before this PR, LTX2TwoStagesPipeline + cache_backend=cache_dit inherited LTX2Pipeline.post_load_weights's cache_dit branch and raised ValueError("Cache-DiT: no enabler registered for pipeline class 'LTX2TwoStagesPipeline'..."). Now the run proceeds at full denoise cost with one log warning. Do you think this works fine? @zhenhuaw-me

I think we should explicitly raise exceptions

luyiyun1021 · 2026-07-21T06:57:01Z

Suggested test additions (all in tests/unittest/_torch/visual_gen/test_cache_dit.py

Compile-on integration test (GPU) — in TestCacheDiTRealPipelineForward, add a Wan 2.1 1.3B variant with torch_compile.enable=True. Assert after PipelineLoader.load: the blocks cache_dit holds are torch._dynamo.OptimizedModule, and after one forward, dynamo counters show compiled frames actually ran.
Why: nothing currently locks the enable-after-compile ordering — every existing cache_dit test runs with compile off, so reverting the order keeps CI green, and the regression is perf-only (~25% on Wan) and invisible to correctness assertions. The counters check is needed because OptimizedModule blocks can still silently fall back to eager (LTX-2 does today).
Flux enabler check-flags unit test (CPU) — new test class. Call enable_cache_dit_for_flux with blocks wrapped by real torch.compile (lazy, no actual compilation) vs. plain eager blocks, patch BlockAdapter to capture kwargs, and assert check_forward_pattern/check_num_outputs are False for compiled and True for eager.
Why: compiled Flux + cache_dit is a supported config (benchmarked in this PR) with zero test coverage; if _has_compiled_blocks breaks or the flags are inverted, the server crashes at startup on cache_dit's inspect assert while CI stays green.
Registry ↔ declaration consistency test (CPU) — new test class. Assert bidirectionally that every pipeline class in CUSTOM_CACHE_DIT_ENABLERS declares "cache_dit" in SUPPORTED_CACHE_BACKENDS, and vice versa.
Why: the new opt-in gate turns a missing declaration into a log warning instead of an error, so a pipeline wired to an enabler but missing the declaration silently loses cache acceleration. The Qwen-Image reapply ([None][feat] Qwen-Image TeaCache/Cache-DiT Support #16339) is the first realistic case.

o-stoner

Left a comment on perhaps enforcing explicit exception raising rather than logging, but otherwise LGTM, thanks!

chang-l · 2026-07-23T21:22:14Z

/bot run

tensorrt-cicd · 2026-07-23T21:28:56Z

PR_Github #61405 [ run ] triggered by Bot. Commit: 19c8e02 Link to invocation

tensorrt-cicd · 2026-07-23T21:48:44Z

PR_Github #61405 [ run ] completed with state FAILURE. Commit: 19c8e02
/LLM/main/L0_MergeRequest_PR pipeline #49632 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chang-l · 2026-07-23T22:28:31Z

/bot run --disable-fail-fast

Signed-off-by: BrianLi23 <briannli233@gmail.com>

tensorrt-cicd · 2026-07-23T22:35:04Z

PR_Github #61417 [ run ] triggered by Bot. Commit: 341edf4 Link to invocation

Signed-off-by: BrianLi23 <briannli233@gmail.com>

chang-l · 2026-07-24T00:17:48Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-24T00:24:21Z

PR_Github #61432 [ run ] triggered by Bot. Commit: e3d43bd Link to invocation

tensorrt-cicd · 2026-07-24T00:26:16Z

PR_Github #61417 [ run ] completed with state ABORTED. Commit: 341edf4

Link to invocation

Signed-off-by: BrianLi23 <briannli233@gmail.com>

luyiyun1021 · 2026-07-24T03:40:51Z

LGTM. Thanks.

tensorrt-cicd · 2026-07-24T03:46:34Z

PR_Github #61432 [ run ] completed with state SUCCESS. Commit: e3d43bd
/LLM/main/L0_MergeRequest_PR pipeline #49657 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

chang-l · 2026-07-24T03:49:26Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-24T03:55:51Z

PR_Github #61493 [ run ] triggered by Bot. Commit: ed0d8b6 Link to invocation

tensorrt-cicd · 2026-07-24T07:42:18Z

PR_Github #61493 [ run ] completed with state SUCCESS. Commit: ed0d8b6
/LLM/main/L0_MergeRequest_PR pipeline #49712 completed with status: 'SUCCESS'

CI Report

Link to invocation

reorder cache-dit and torch.compile application order

40a7b4c

Signed-off-by: BrianLi23 <briannli233@gmail.com>

github-actions Bot assigned BrianLi23 Jul 16, 2026

github-actions Bot added the VisualGen label Jul 16, 2026

BrianLi23 added 3 commits July 16, 2026 11:34

remove comments

f4dd432

Signed-off-by: BrianLi23 <briannli233@gmail.com>

remove stale test

aa9ce3e

Signed-off-by: BrianLi23 <briannli233@gmail.com>

Merge branch 'main' into brian-li-cache-dit-torch-compile-resolve

d562d09

BrianLi23 marked this pull request as ready for review July 16, 2026 21:29

BrianLi23 requested a review from a team as a code owner July 16, 2026 21:29

BrianLi23 requested review from karljang and luyiyun1021 July 16, 2026 21:29

coderabbitai Bot reviewed Jul 16, 2026

View reviewed changes

Comment thread tests/unittest/_torch/visual_gen/test_teacache.py

handle cache support per pipeline

19c8e02

Signed-off-by: BrianLi23 <briannli233@gmail.com>

zhenhuaw-me requested a review from o-stoner July 17, 2026 01:30

zhenhuaw-me approved these changes Jul 17, 2026

View reviewed changes

o-stoner reviewed Jul 21, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/pipeline.py Outdated

chang-l enabled auto-merge (squash) July 23, 2026 21:22

remove supported cache backend

341edf4

Signed-off-by: BrianLi23 <briannli233@gmail.com>

auto-merge was automatically disabled July 23, 2026 22:30
Head branch was pushed to by a user without write access

add additional tests

e3d43bd

Signed-off-by: BrianLi23 <briannli233@gmail.com>

delete unneeded test

ed0d8b6

Signed-off-by: BrianLi23 <briannli233@gmail.com>

luyiyun1021 approved these changes Jul 24, 2026

View reviewed changes

zhenhuaw-me merged commit a8b5409 into NVIDIA:main Jul 24, 2026
7 checks passed

Uh oh!

Conversation

BrianLi23 commented Jul 16, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test Coverage

Follow up

Summary by CodeRabbit

Uh oh!

BrianLi23 commented Jul 16, 2026

Uh oh!

github-actions Bot commented Jul 16, 2026

GitHub Bot Help

kill

skip

reuse-pipeline

Uh oh!

BrianLi23 commented Jul 16, 2026

Uh oh!

coderabbitai Bot commented Jul 16, 2026

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zhenhuaw-me commented Jul 17, 2026

Uh oh!

zhenhuaw-me left a comment

Choose a reason for hiding this comment

Uh oh!

zhenhuaw-me commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

luyiyun1021 commented Jul 21, 2026

Uh oh!

luyiyun1021 commented Jul 21, 2026

Uh oh!

zhenhuaw-me commented Jul 21, 2026

Uh oh!

luyiyun1021 commented Jul 21, 2026

Uh oh!

o-stoner left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

chang-l commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

chang-l commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

chang-l commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

luyiyun1021 commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

chang-l commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

Uh oh!

Reviewers

BrianLi23 commented Jul 16, 2026 •

edited by coderabbitai Bot

Loading